# -*- coding: utf-8 -*-
# May Saint Diana bless your coding!
# Created Time : 2021/12/2 21:50
# Author:Zhou Chenghao

import numpy as np
import torch
from torch.autograd import Variable
import random
import os
import matplotlib.pyplot as plt
from ANC_tools import ANR, insertToHead, insertToTail
import torch.nn as nn

"""
不是构造一个个单独的实验数据样本，而是像经典的FXLMS一样实现一个跟随时域过程进行的迭代
"""
# 读取实验数据
path = "非线性1.csv"  # 一个参考信号和该信号经过非线性过程的对应输出
data = np.loadtxt(path, delimiter=",")

#  自定义一个全连接网络 32*64*16*1  当前网络的输出和之前32个时间点的输入有关
net = torch.nn.Sequential(
    torch.nn.Linear(32, 128),
    torch.nn.ReLU(),
    torch.nn.Linear(128, 32),
    torch.nn.ReLU(),
    torch.nn.Linear(32, 1))

# 如果存在训练过的网络参数文件，直接读取
if os.path.exists('bp_anc2_net_params.pkl'):
    net.load_state_dict(torch.load('bp_anc2_net_params.pkl'))  # 提取net的参数

# for _, param in enumerate(net.parameters()):
#     print(param)
#     print('----------------')


class My_loss(nn.Module):
    """
    自定义损失函数
    """

    def __init__(self):
        super().__init__()
        self.spm = torch.tensor([0.1, 0, 0, 0.6, 0.3, 0.2, -0.1, 0.08])

    def forward(self, x, y):
        """
        网络得到的输出要继续经过次级路径，得到降噪点处的信号
        :param x:
        :param y:
        :return:
        """
        global tensor_cy
        tensor_cy[0] = x
        filtered_y = torch.matmul(tensor_cy, self.spm)
        return torch.pow(torch.abs(filtered_y + y), 2)


optimizer = torch.optim.SGD(net.parameters(), lr=0.1)  # 用SGD优化器优化，lr为学习率，将net的参数传入
loss_func = My_loss()
# loss_func = torch.nn.MSELoss()
torch.autograd.set_detect_anomaly(True)
# plt.ion()  # 画图
# plt.show()
# 构造简单训练数据
secondaryPathModel = np.array([0.1, 0, 0, 0.6, 0.3, 0.2, -0.1, 0.08])

signal_x = data[:, 0]
signal_y = data[:, 1]

T = len(signal_x)
cx = np.zeros(32)
cy = np.zeros(8)
error = np.zeros(T)
ys = np.zeros(T)
tensor_cy = torch.zeros(8)

print("----------开始训练----------")
for i in range(1000):
    # 更新参考信号
    cx = insertToHead(cx, signal_x[i])
    tensor_cx = Variable(torch.from_numpy(cx)).to(torch.float)
    # 通过网络计算得到的预测值
    prediction = net(tensor_cx)

    # # 网络的输出传递给次级扬声器
    cy = insertToHead(cy, prediction.data.numpy())
    # # 次级噪声是次级扬声器发出的信号经过次级路径的结果
    ys[i] = np.dot(cy, secondaryPathModel)
    # # 得到误差信号
    error[i] = ys[i] + signal_y[i]
    # # 更新网络
    # # 计算ys和yp的误差，得到损失函数

    tensor_cy = torch.from_numpy(cy).float()
    tensor_yp = torch.from_numpy(np.array(signal_y[i])).float()

    loss = loss_func(prediction, tensor_yp)
    loss.requires_grad_(True)

    optimizer.zero_grad()  # 将net.parameters中上一批梯度清零
    # print(loss.data, loss.grad, loss.grad_fn)
    loss.backward()  # 反向传递  计算节点梯度  retain_graph=True

    optimizer.step()  # 优化节点
    # if t % 5 == 0:  # 动态显示
    #     # plot and show learning process
    #     plt.cla()
    #     plt.plot(y.data.numpy(), lw=5)
    #     prediction_output = np.dot(prediction.data.numpy(), secondaryPathModel.t().data.numpy())
    #     plt.plot(prediction_output, 'r-', lw=5)
    #     plt.text(25, 3, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color': 'green'})
    #     plt.pause(0.1)
print("----------训练完成！----------")


# 保存网络的参数
torch.save(net.state_dict(), 'bp_anc2_net_params.pkl')  # 只保存网络中的参数 (速度快, 占内存少)
plt.figure()


plt.plot(signal_y)
plt.plot(ys)
plt.plot(error)
plt.legend(['primary noise', 'secondary noise', 'residual noise'])
plt.show()

plt.figure()
plt.plot(ANR(signal_y, error))
plt.title('BP_ANR')
plt.show()

# for _, param in enumerate(net.parameters()):
#     print(param)
#     print('----------------')
# plt.ioff()
# plt.show()

# 验证拟合效果
#
# cx = np.zeros(32)
# cy = np.zeros(8)
# BP_error = np.zeros(len(signal_x))
# print("----------验证拟合效果----------")
# for i in range(len(signal_x)):
#     cx = insertToTail(cx, signal_x[i])  # 需要和网络的数据对应 从上至下应该是 x(n-15) 到 x(n)，所以要插入到末尾
#     BP_x = Variable(torch.from_numpy(cx)).to(torch.float)
#     BP_pred = net(BP_x)  # 得到交给次级扬声器的输出
#     cy = insertToHead(cy, BP_pred.data.numpy())  # 更新y
#     filtered_y = np.dot(secondaryPathModel.data.numpy(), cy[:, np.newaxis])
#     BP_error[i] = filtered_y + signal_y[i]
# print("平均误差：{0},最大误差：{1}".format(np.average(np.abs(BP_error)),
#                                  np.max(np.abs(BP_error))))
#
# plt.subplot(121)
# plt.plot(signal_y)
# plt.plot(BP_error)
# plt.legend(['y', 'error'])
# plt.title('BP')
# plt.subplot(122)
# plt.plot(ANR(signal_y, BP_error))
# plt.title('BP_ANR')
# plt.show()
"""
尝试模仿真正的物理系统，设计思路和遇到的问题总结如下：
1、设计思路：网络的更新过程像自适应滤波一样，跟随时域的前进而进行
    这样能保存之前的输出值，也就可以进行次级路径的输出
2、问题：次级路径模型里面，x(n)项的系数基本是0，那么在求梯度的时候，就会消失，也就无法更新网络了
    
"""